
import os
import re
import sys
import glob
import json
import threading
import multiprocessing

import cv2
import numpy

from Module.xRayDetection.xray.logs import Logger
from Module.xRayDetection.xray.mes import request, prettify
from Module.xRayDetection.xray.util import stringify, get_devices
from Module.xRayDetection.xray.ftp import FTP, put_queue, save_cache, check_folder

if os.name == 'nt':
    import pathlib
    pathlib.PosixPath = pathlib.WindowsPath


lines = {
    'c23_cv04': {
        'sep': '_',
        'extra': True,
        'host': 'czccd.svolt.cn',
        'user': 'CM02YCXR01',
        'passwd': 'SVOLT@2021',
        'prefix':'/CZ-PH2-ZP-L2-YCRAY00',
        'suffix': '-originalimgs',
    },
    'c21_cm12': {
        'sep': '-',
        'extra': False,
        'host': 'czccd.svolt.cn',
        'user': 'CM01YCXR01',
        'passwd': 'SVOFTPCM01YCXR01',
        'prefix':'/CM01YCXR01/ZY-XRay/CZ-PH2-ZP-L1-YCXRAY00',
        'suffix': '-originalimgs/Image/SRC',
    },
    'c31_cv04': {
        'sep': '-',
        'extra': False,
        'host': 'czccd.svolt.cn',
        'user': 'CV04YCXR',
        'passwd': 'Svolt@2022',
        'prefix':'/CV04YCXR/L1/CZ3CL1XRAYJ0',
        'suffix': '/SRC',
    }
}


def remove_colors(im, box=None, thres=20, maxval=255):

    roi = im[int(box[1]):int(box[3]), int(box[0]):int(box[2]), ...] if box is not None else im

    if roi is None:
        return None

    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)

    colors = cv2.threshold(hsv[:,:,1:2], thres, maxval, cv2.THRESH_BINARY)

    dst = cv2.inpaint(roi, colors[1], 1, cv2.INPAINT_NS)

    if box is not None:
        im[int(box[1]):int(box[3]), int(box[0]):int(box[2]), ...] = dst
    else:
        im = dst
    return im


def get_region(i, im):
    im = numpy.array(im)
    if i == 0:
        xmin = 0
        xmax = im.shape[1] // 2
        ymin = 0
        ymax = im.shape[0] // 2
    if i == 1:
        xmin = im.shape[1] // 2
        xmax = im.shape[1]
        ymin = 0
        ymax = im.shape[0] // 2
    if i == 2:
        xmin = 0
        xmax = im.shape[1] // 2
        ymin = im.shape[0] // 2
        ymax = im.shape[0]
    if i == 3:
        xmin = im.shape[1] // 2
        xmax = im.shape[1]
        ymin = im.shape[0] // 2
        ymax = im.shape[0]

    return xmin, xmax, ymin, ymax


def draw(img, det):
    img = numpy.array(img)
    for i, box in enumerate(det):
        cv2.rectangle(img,
            (int(box[0]), int(box[1])), (int(box[2]), int(box[3])),
            (0,0,255), 2)
        cv2.putText(img, str(i), (int(box[0]), int(box[1])),
            cv2.FONT_HERSHEY_PLAIN, img.shape[0] // 1000, (0,0,255), img.shape[0] // 1000)
        cv2.putText(img, f'{i}: d{int(box[4] * 100)}' + (f' c{int(box[6] * 100)}' if len(box) > 6 else ''),
            (img.shape[1] // 10 * 8, img.shape[1] // 20 * (i+1)),
            cv2.FONT_HERSHEY_PLAIN, img.shape[0] // 500, (0,0,255), img.shape[0] // 500)

    cv2.putText(img, 'NG' if len(det) > 0 else 'OK', (img.shape[1] // 50, img.shape[0] // 20),
        cv2.FONT_HERSHEY_PLAIN, img.shape[0] // 250, ((0,0,255) if len(det) > 0 else (0,255,0)),
        img.shape[0] // 200)

    return img


def tolabmelme(path:str, img, dets, model=None, names=None):
    '''
    <annotation>
        <filename>1.jpg</filename>
        <size>
            <height></height>
            <width></width>
            <depth></depth>
        </size>
        <object>
            <bndbox>
                <xmin></xmin>
                <ymin></ymin>
                <xmax></xmax>
                <ymax></ymax>
            </bndbox>
            <name></name>
        </object>
    </annotation>
    '''
    img = numpy.array(img)
    labelme = f'<annotation>\
                    <filename>{os.path.basename(path)}</filename>\
                    <weights>{os.path.basename(model.model.w) if model else ""}</weights>\
                    <size>\
                        <height>{img.shape[0]}</height>\
                        <width>{img.shape[1]}</width>\
                        <depth>{img.shape[2]}</depth>\
                    </size>'
    def to_numpy(n) -> numpy.ndarray:
        return numpy.array(n)
    for det in dets:
        if len(det) < 7:
            det = to_numpy(det).tolist() + [0] * 9
        labelme += f'\
        <object>\
            <bndbox>\
                <xmin>{int(det[0])}</xmin>\
                <ymin>{int(det[1])}</ymin>\
                <xmax>{int(det[2])}</xmax>\
                <ymax>{int(det[3])}</ymax>\
            </bndbox>\
            <score>\
                <detection>{det[4]}</detection>\
                <classification>{det[6] if type(det[6]) is not str else 1}</classification>\
            </score>\
            <name>{model.model.names[int(det[5])] if model else names[int(det[5])] or ""}</name>\
        </object>'
    labelme += '</annotation>'
    return labelme


def get_lines():
    return list(lines.keys())


class XRay:

    def __init__(
            self,
            host='',
            line='',
            pc=0,
            device='',
            ftp=None,
            logger=None,
            access_token='',
            content='',
            interval=60,
            mes=False,
            local_img=True,
            remote_img=True,
            local_xml=False,
            remote_xml=False,
            split=False,
            augment=False,
            loop=True,
            sub_queue=None,
            main_queue=None,
            sep='',
            cache_path='',
            save_dir='',
            remote_dir='',
            local_model_dir='',
            upload_dir='',
            remote_model_dir='',
            git_dir='',
            roi_model='',
            cls_model='',
            fold_model='',
            roi_weights='',
            cls_weights='',
            fold_weights='',
            roi_thres=0.5,
            cls_thres=0.5,
            fold_thres=0.5,
            shape=1280,
            crop=False,
            folder=''
    ):

        user = lines[line]['user']
        passwd = lines[line]['passwd']
        self.host = host or lines[line]['host']
        root = os.path.dirname(__file__)

        self.pc = pc
        self.line = line

        self.mes = mes
        self.loop = loop
        self.crop = crop
        self.split = split
        self.shape = shape
        self.folder = folder
        self.local_img = local_img
        self.local_xml = local_xml
        self.remote_img = remote_img
        self.remote_xml = remote_xml

        self.sub_queue = sub_queue
        self.main_queue = main_queue

        self.sep = sep or lines[line]['sep']
        self.extra = lines[line]['extra']
        self.suffix = lines[line]['suffix']
        self.prefix = lines[line]['prefix']

        self.content = content
        self.interval = interval
        self.access_token = access_token

        self.save_dir = save_dir or os.path.join(root, 'downloads')
        self.upload_dir = upload_dir or '/'.join(['', line, f'#{pc}']) if upload_dir is not None else upload_dir
        self.remote_dir = remote_dir or f'{self.prefix}{pc}{self.suffix}'
        self.cache_dir = os.path.dirname(cache_path) if os.path.exists(cache_path) else os.path.join(self.save_dir, line, f'#{pc}' if pc else '')
        self.local_model_dir = local_model_dir or os.path.join(root, 'downloads', 'models')
        self.logger = logger or Logger(os.path.join(self.save_dir, 'logs'))
        self.ftp = ftp or FTP(host=self.host if not os.path.exists(remote_dir) else '', user=user, passwd=passwd, logger=self.logger)

        self.fold_weights = fold_weights
        self.roi_weights = roi_weights
        self.cls_weights = cls_weights

        self.cls_thres = float(cls_thres)
        self.roi_thres = float(roi_thres)
        self.fold_thres = float(fold_thres)

        device = device or [d.split(':')[-1] for d in get_devices(device)][-1]

        self.download_models(remote_model_dir)

        from Module.xRayDetection.xray.yolov5 import YOLOV5
        self.fold_model = fold_model or YOLOV5(
            weights=self.fold_weights, 
            device=device, 
            logger=self.logger, 
            git_dir=git_dir, 
            augment=augment) \
            if self.fold_weights and float(fold_thres) > 0 else None
        self.cls_model = cls_model or YOLOV5(
            weights=self.cls_weights, 
            device=device, 
            logger=self.logger, 
            git_dir=git_dir, 
            augment=augment) \
            if self.cls_weights and float(cls_thres) > 0 else None
        self.roi_model = roi_model or YOLOV5(
            weights=self.roi_weights, 
            device=device, 
            logger=self.logger, 
            git_dir=git_dir, 
            augment=augment) \
            if self.roi_weights and float(roi_thres) > 0 else None
        
        cfg_ = stringify(self.__dict__)
        self.logger.log(self.logger.INFO, json.dumps(cfg_, indent=4))

        self.cache = cfg_


    def download_models(self, remote_model_dir=''):

        if remote_model_dir:
            ftp = FTP(
                host=self.host or lines[get_lines()[0]]['host'],
                user=lines[get_lines()[0]]['user'],
                passwd=lines[get_lines()[0]]['passwd'],
                logger=self.logger
            )
            for remote_model in self.ftp.listdir(remote_model_dir):
                if self.line in remote_model['name'] and \
                    not os.path.exists(os.path.join(self.local_model_dir, remote_model['name'])):
                    put_queue(self.sub_queue, {'msg': f"Start downloading {remote_model['path']}"})
                    ftp.download(download_path=remote_model['path'], path=os.path.join(self.local_model_dir, remote_model['name']))
                    put_queue(self.sub_queue, {'msg': f"Finshing downloading {remote_model['path']}"})


        for type_ in ['fold', 'cls', 'roi']:
            weights = sorted(glob.glob(os.path.join(self.local_model_dir, self.line + f'_{type_}*.pt')))
            try:
                setattr(self, f'{type_}_weights', getattr(self, f'{type_}_weights') or weights[-1])
            except Exception as e:
                self.logger.log(self.logger.ERROR, f"{e} No {type_}_weights for {self.line}", exception=True)

            for weight in weights[:-1]:
                put_queue(self.sub_queue, {'msg': f'Remove obsolete model {weight}'})
                self.logger.log(self.logger.INFO, f'Remove obsolete model {weight}')
                os.remove(weight)


    def get_dst(self, path:str, dst, ext=''):
        dst_ = path.replace(self.cache_dir, dst)
        dst_ = dst_.replace(self.remote_dir, dst)
        dst_ = dst_.replace(self.prefix, f'{dst}/#')
        dst_ = dst_.replace(self.suffix, '').replace('\\', '/')
        dst_ = dst_.replace(path.split('.')[-1], ext) if ext else dst_
        return dst_


    def get_barcode(self, path:str):
        barcode = os.path.basename(path).split(self.sep)[0]
        try:
            path_ = path[:path.index(barcode)+len(barcode)] + ('_' + os.path.basename(path).split(self.sep)[-2] if self.extra else '')
        except Exception as e:
            msg = f'can not parse {self.line} barcode from {path}'
            put_queue(self.sub_queue, {'msg': msg})
            self.logger.log(self.logger.ERROR, f'{e} {msg}', exception=True)
            path_ = ''
        return path_


    def detect_fold(self, im, roi:list=[], regions=[]):

        dets = [[], [], [], []]

        if not self.merge:

            im, det = self.fold_model.detect(im, self.fold_thres, roi)

            return im, det

        for i in range(4):

            xmin, xmax, ymin, ymax = get_region(i, im)

            img = im[ymin:ymax, xmin:xmax, ...]
            det = []

            is_roi = len(roi) < 2 and len([r for r in roi if len(r) < 5])

            if is_roi:
                box_ = [box for box in roi.copy()
                    if box[0] > xmin - xmin / 99 and box[0] < xmax 
                    and box[1] > ymin - ymin / 99 and box[1] < ymax
                    and box[2] > xmin - xmin / 99 and box[2] < xmax 
                    and box[3] > ymin - ymin / 99 and box[3] < ymax]

                if len(box_) == 0:
                    continue

                box = box_[0]
                box[:4] = [max(b, 0) for b in (box[:4] - [xmin, ymin, xmin, ymin])]
            else:
                box = []

            img = remove_colors(img, box=box) if is_roi else img
            im[ymin:ymax, xmin:xmax, ...] = img

            for region in regions:
                im[region[1]+ymin:region[3]+ymin,region[0]+xmin:region[2]+xmin,...] = remove_colors(im[region[1]+ymin:region[3]+ymin,region[0]+xmin:region[2]+xmin,...])

            _, det = self.fold_model.detect(img, self.fold_thres, box)

            if len(det) == 0:
                continue

            try:
                import torch
                import torchvision

                det_ = []
                roi_ = torch.Tensor(roi or [[]])[:, :4]
                for d in det:
                    if len(roi) == 0 or torchvision.ops.box_iou(roi_, torch.Tensor(d)[:4] + torch.Tensor([xmin,ymin]).repeat(1,2)).max() > 0:
                        det_.append(d)
                dets[i] = det_
            except:
                pass

        return im, dets


    def classify_fold(self, im, det, dets:list):
        im_ = im[int(det[1]):int(det[3]), int(det[0]):int(det[2]), ...]
        pred = self.cls_model.classify(im_)
        if pred[0] > self.cls_thres:
            dets.append(list(det) + [pred[0]])
        return dets


    def filter_fold(self, im, dets):
        if self.cls_model is None:
            return dets

        if len(dets) and type(dets[0]) is list:
            dets_ = [[],[],[],[]]
            for i, det_ in enumerate(dets):
                for det in det_:
                    xmin, xmax, ymin, ymax = get_region(i, im)
                    img = im[ymin:ymax, xmin:xmax, ...]
                    dets_[i] = self.classify_fold(img, det, dets_[i])
        else:
            dets_ = []
            for det in dets:
                dets_ = self.classify_fold(im, det, dets_)
        return dets_


    def save_result(self, det, im, im_path, xml_path:str):

        if self.local_img:
            os.makedirs(os.path.dirname(im_path), exist_ok=True)
            cv2.imwrite(im_path, im)

        if self.local_xml:
            xml_  = tolabmelme(im_path, im, det, self.fold_model)
            xml_ = prettify(xml_)

            os.makedirs(os.path.dirname(xml_path), exist_ok=True)
            open(xml_path.replace('.xml', f'-{os.path.basename(self.fold_model.model.w)}.xml'), 
                'w', encoding='utf-8').write(xml_)
            return xml_


    def post_processing(self, results):

        try:

            path = self.get_barcode(results[0]['path']) + '.jpg'

            for i in range(4):
                try:
                    imsz = numpy.array(results[i]['im']).shape
                    break
                except:
                    pass

            im = numpy.zeros((imsz[1]*2, imsz[0]*2, 3))

            det = []
            for pos, result in enumerate(results):
                xmin, xmax, ymin, ymax = get_region(pos, im)

                if hasattr(result['im'], 'shape'):
                    im[ymin:ymax, xmin:xmax, ...] = numpy.array(result['im']).copy()

                det += [(numpy.array(d[:4]).astype(int) + [xmin,ymin,xmin,ymin]).tolist() + \
                        [d[-2], int(d[-1]), self.fold_model.model.names[d[-1]]] for d in result['det']]

            if len(det) == 0:
                return im, det

            xml_ = self.save_result(det, im,
                im_path=self.get_dst(path, self.cache_dir, 'jpg'),
                xml_path=self.get_dst(path, self.cache_dir, 'xml')
            )

            if self.upload_dir and self.ftp:

                if self.remote_img:
                    self.ftp.upload(
                        upload_path=self.get_dst(path, self.upload_dir, 'jpg'),
                        bytes_=cv2.imencode('.jpg', im)[1]
                    )

                if self.remote_xml:
                    self.ftp.upload(
                        upload_path=self.get_dst(path, self.upload_dir, 'xml'), 
                        bytes_=xml_.encode()
                    )

            if self.mes:
                response = request(barcode=[os.path.basename(path).split('_')[0]], logger=self.logger)
                self.logger.log(self.logger.INFO, response)

            return im, det

        except Exception as e:
            self.logger.log(self.logger.ERROR, f'{e} {path}', exception=True)
            return None, None


    def merge_processing(self, result):

        xml_tmp = self.local_xml
        self.local_xml = True
        try:
            path = result['path']
            im_path_ = self.get_dst(path, os.path.join(self.save_dir, 'images'), 'jpg')
            xml_path_ = self.get_dst(path, os.path.join(self.save_dir, 'labels'), 'xml')
            merge_path = self.get_dst(path, os.path.join(self.save_dir, 'merge'), 'jpg')

            os.makedirs(os.path.dirname(merge_path), exist_ok=True)
            cv2.imwrite(merge_path, result['im'])

            for pos, det in enumerate(result['det']):
                if len(det) < 1:
                    continue

                xmin, xmax, ymin, ymax = get_region(pos, result['im'])

                self.save_result(det,
                    im=result['im'][ymin:ymax, xmin:xmax, ...],
                    im_path=im_path_.replace('.jpg', f'_{pos}.jpg'),
                    xml_path=xml_path_.replace('.xml', f'_{pos}.xml')
                )

                if self.crop:
                    fold_path = self.get_dst(path, os.path.join(self.save_dir, 'fold'), 'jpg')
                    for i, d in enumerate(det):
                        os.makedirs(os.path.dirname(fold_path), exist_ok=True)
                        im = result['im'][int(d[1]):int(d[3]), int(d[0]):int(d[2]), ...]
                        cv2.imwrite(fold_path.replace('.jpg', f'_f{i}.jpg'), im)

            self.local_xml = xml_tmp
            return im, det

        except Exception as e:
            self.logger.log(self.logger.ERROR, f'{e} {path}', exception=True)
            return None, None


    def split_processing(self, result):

        xml_tmp = self.local_xml
        self.local_xml = True
        try:
            path = result['path']
            im = result['im']
            det = result['det']

            self.save_result(det,
                im=im,
                im_path=self.get_dst(path, os.path.join(self.save_dir, 'images'), 'jpg'),
                xml_path=self.get_dst(path, os.path.join(self.save_dir, 'labels'), 'xml')
            )

            self.local_xml = xml_tmp
            return im, det

        except Exception as e:
            self.logger.log(self.logger.ERROR, f'{e} {path}', exception=True)
            return None, None


    def is_barcode_equal(self, path1, path2):
        return self.get_barcode(path1) == self.get_barcode(path2) and len(path1) == len(path2)


    def process_xray(self, path='', i=0, len_files=0, files=[], results:list=[], obj:dict={}, shape=1280, **kwargs):

        dir_ = path.replace(path.split('/')[-1], '')[:-1]

        obj.update({
            'im': None,
            'det': None,
            'msg': f'XRay start: {path}',
        })

        put_queue(self.sub_queue, obj)

        bio = self.ftp.download(download_path=path) if not os.path.exists(path) else open(path, 'rb').read()

        im = numpy.asarray(bytearray(bio.getvalue() if not os.path.exists(path) else bio), dtype='uint8')

        try:
            im = cv2.imdecode(im, cv2.IMREAD_COLOR)
            self.merge = numpy.array(im).shape[0] == shape
        except Exception as e:
            self.logger.log(self.logger.ERROR, e, exception=True)
            print(kwargs)
            return results

        c = 0
        roi = []

        while self.roi_model and len(roi) == 0 and c < 9:
            im = cv2.rotate(im, cv2.ROTATE_90_COUNTERCLOCKWISE) if not self.merge else im
            _, roi = self.roi_model.detect(im, self.roi_thres) if self.roi_model else []
            c += 1

        self.logger.log(self.logger.INFO, 'roi ' + '\n'.join([str(r) for r in roi]))

        im_, det = self.detect_fold(im, roi)

        det = self.filter_fold(im, det) if self.cls_model else det

        self.logger.log(self.logger.INFO, 'fold ' + '\n'.join([str(d) for d in det]))

        result = {
            'im': im_ if len(roi) else im,
            'roi': roi,
            'det': det,
            'path': path,
        }

        ims = det_ = None

        if not self.merge and not self.split:

            files_ = [f for f in files[max(i - 9, 0):i + 9] if self.is_barcode_equal(f, path)]

            def key(f:str):
                return f.split(self.sep)[-2]

            temp = sorted(files_, key=key)[:4]
            files_ = temp if not self.extra and path in temp else files_
            temp = sorted(files_, key=key)[4:]
            files_ = temp if not self.extra and path in temp else files_

            if files_.index(path) >= 0 and files_.index(path) < len(files_) and (len([r for r in results if r['path'] == result['path']]) == 0):
                results.append(result)

            results_ = [r for r in results if self.is_barcode_equal(r['path'], path)]

            if len(results_) == len(files_):

                if hasattr(sys, 'gettrace') and sys.gettrace() is not None:
                    print(f"{str(i+1).rjust(len(str(len_files)),'0')}/{len_files}")
                else:
                    self.logger.log(self.logger.INFO, f"{str(i+1).rjust(len(str(len_files)),'0')}/{len_files}")

                if len(results_) != 4 or (len(results_) == 4 and \
                    len([r for r in results_ if len(r['det']) > 0 or len(r['im']) == 0]) > 0):

                    ims, det_ = self.post_processing(results_)
                    self.cache[dir_] = f'{i + 1}/{str(len_files)}'
                    save_cache(cache=self.cache, local_cache_path=self.cache['local_cache_path'], remote_cache_path=self.cache['remote_cache_path'], ftp=self.ftp)

                elif self.sub_queue:
                    ims, det_ = self.post_processing(results_)

                results = [r for r in results if len([r_ for r_ in results_ if r['path'] == r_['path']]) == 0]
        else:
            if hasattr(sys, 'gettrace') and sys.gettrace() is not None:
                print(f"{str(i+1).rjust(len(str(len_files)),'0')}/{len_files}")
            else:
                self.logger.log(self.logger.INFO, f"{str(i+1).rjust(len(str(len_files)),'0')}/{len_files}")

            if len([d for d in det if len(d)]):
                ims, det_ = self.split_processing(result) if self.split else self.merge_processing(result)
                self.cache[dir_] = f'{i + 1}/{str(len_files)}'
                save_cache(cache=self.cache, local_cache_path=self.cache['local_cache_path'], remote_cache_path=self.cache['remote_cache_path'], ftp=self.ftp)

            elif self.sub_queue:
                ims, det = self.split_processing(result) if self.split else self.merge_processing(result)

        obj.update({
            'im': ims,
            'det': det_,
            'msg': f'XRay finish: {path}',
        })

        put_queue(self.sub_queue, obj)

        if i + 1 == len_files:
            self.cache[dir_] = f'{i + 1}/{str(len_files)}'
            save_cache(cache=self.cache, local_cache_path=self.cache['local_cache_path'], remote_cache_path=self.cache['remote_cache_path'], ftp=self.ftp)

        return results


    def run(self):

        if self.loop and self.access_token:
            def content(remote, len_files):
                return f"Uploading stopped. {self.line} {self.pc} {str(remote).replace(self.remote_dir, '')} {len_files}"

            kwargs = dict(
                ftp=self.ftp,
                func=content,
                access_token=self.access_token,
                cache_path=os.path.join(self.cache_dir, 'stop.txt'),
                remote_dir=self.remote_dir,
                interval=self.interval
            )

            threading.Thread(target=check_folder, kwargs=kwargs).start()

        self.ftp.loop_func(
            self.process_xray, loop=self.loop, dir_=self.save_dir, items=[self.remote_dir], cache_dir=self.cache_dir or self.upload_dir,
            main_queue=self.main_queue, sub_queue=self.sub_queue, cache=self.cache, shape=self.shape, folder=self.folder)

        del self.fold_model
        del self.cls_model
        del self.roi_model
        import gc
        import torch
        gc.collect()
        torch.cuda.empty_cache()


def run(**kwargs:dict):
    if 'files' in kwargs:
        files = kwargs.pop('files', [])
        len_files = len(files)
        xray = XRay(**kwargs)
        progress = str(xray.cache[[*xray.cache][0]]).split('/') if len([*xray.cache]) else [0,0]
        start = 0 if len(progress) < 2 or int(progress[1]) != len_files else int(progress[0])
        xray.cache['local_cache_path'] = kwargs['cache_path']
        xray.cache['remote_cache_path'] = ''
        for i, path in enumerate(files[start:]):
            if i <= start:
                i = start
            xray.process_xray(path=path, i=i, len_files=len_files, shape=kwargs.pop('shape', xray.shape), **kwargs)
        try:
            xray.cache[[*xray.cache][0]] = f'{len_files}/{len_files}'
            json.dump(xray.cache, open(kwargs['cache_path'], 'w', encoding='utf-8'))
        except:
            pass
    else:
        xray = XRay(**kwargs)
        xray.run()

def multi(**kwargs):

    import tqdm
    import torch
    split = kwargs.pop('split', torch.cuda.device_count())
    cache_path = kwargs.pop('cache_path', 0)
    codes = kwargs.pop('codes', '')

    files = kwargs.pop('files', None)

    if files:
        files_ = files
    elif not kwargs.pop('filter', '') and os.path.exists(os.path.join(kwargs['save_dir'], kwargs.get('name', 'filter')+'.txt')):
        files_ = open(os.path.join(kwargs['save_dir'], kwargs.pop('name', 'filter')+'.txt')).read().split('\n')
        codes_ = [os.path.basename(file).split(kwargs.get('sep', '/'))[0].upper() for file in files_]
    else:
        files = glob.glob(os.path.join(kwargs['remote_dir'], '**', '**.jpg'), recursive=True)

        try:
            codes = open(codes).read().upper()
        except:
            codes = ''
            pass

        codes_ = []
        files_ = []
        
        for file in tqdm.tqdm(files):
            code_ = os.path.basename(file).split(kwargs.get('sep', '_'))[0].upper()
            if code_ in codes:
                codes_.append(code_)
                files_.append(file)
            if not codes:
                files_.append(file)

        os.makedirs(kwargs['save_dir'], exist_ok=True)
        open(os.path.join(kwargs['save_dir'], kwargs.pop('name', 'filter')+'.txt'), 'w').write('\n'.join(files_))

    length = len(files_)

    multiprocessing.set_start_method('spawn')

    for i in range(split):
        kwargs['cache_path'] = str(cache_path).replace('.json', str(i+1).rjust(2,'0') + '.json')
        kwargs['files'] = files_[i*(length // split):(i+1)*(length // split) if (i+2)*(length // split) < length else length]
        kwargs['device'] = str(i % torch.cuda.device_count())
        os.environ['CUDA_VISIBLE_DEVICES'] = kwargs['device']

        if hasattr(sys, 'gettrace') and sys.gettrace() is not None:
            kwargs['interval'] = 1
            run(**kwargs)

        multiprocessing.Process(target=run, kwargs=kwargs).start()

def test(args, **kwargs):

    if hasattr(sys, 'gettrace') and sys.gettrace() is not None:
        kwargs['interval'] = 1
        run(**kwargs)
        args = []

    multiprocessing.set_start_method('spawn')

    for arg in args:
        kwargs['pc'] = arg['pc']
        kwargs['device'] = arg['device']
        os.environ['CUDA_VISIBLE_DEVICES'] = str(arg['device'])
        multiprocessing.Process(target=run, kwargs=kwargs).start()
